package com.starhub.application.rag.service.impl;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.filter.Filter;
import lombok.extern.slf4j.Slf4j;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import com.starhub.application.rag.config.MilvusEmbeddingConfig;

import com.starhub.application.model.biz.VectorModelFactory;
import com.starhub.application.rag.constant.RagConstant;
import com.starhub.application.rag.dto.EmbeddingDto;
import com.starhub.application.rag.service.EmbeddingService;
import com.starhub.application.rag.service.MilvusConfigService;
import com.starhub.application.rag.store.LangchainStoreBuilder;
import com.starhub.application.rag.transformer.SimpleTextSegmentTransformer;
import com.starhub.application.rag.util.DocumentSplitterUtil;
import com.starhub.common.bean.model.MilvusConfig;
import com.starhub.utils.Util;

import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;

import java.util.concurrent.TimeUnit;

/**
 * 向量服务实现类
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-05
 */
@Slf4j
@Service
public class EmbeddingServiceImpl implements EmbeddingService {

    @Autowired
    private MilvusConfigService milvusConfigService;

    @Autowired
    VectorModelFactory vectorModelFactory;

    @Autowired
    SimpleTextSegmentTransformer simpleTextSegmentTransformer;

    @Autowired
    private LangchainStoreBuilder langchainStoreBuilder;

    @Autowired
    private MilvusEmbeddingConfig milvusEmbeddingConfig;

    /**
     * 对输入内容进行向量化
     * @param vectorModel 向量模型
     * @param content 输入内容
     * @return 向量化结果
     * @throws Exception 异常
     */
    public Embedding embedding(String vectorModel, String content) throws Exception{
        // 模型对象实例
        EmbeddingModel embeddingModel = vectorModelFactory.getVectorModel(vectorModel);
        // 对输入内容进行向量化
        Embedding queryEmbedding = embeddingModel.embed(content).content();
        return queryEmbedding;
    }
    
    /**
     * 知识条目向量化 所有的向量化都走这里，只是向量化的内容有区别
     * 
     * @param EmbeddingDto
     * @return 是否成功
     */
    @Override
    public void embedding(EmbeddingDto embeddingDto) throws Exception {
        // 1. 拼接向量化数据内容
        Metadata metadata = new Metadata();
        metadata.put(RagConstant.EmbeddingMetaData.TEXT_ID, Util.null2String(embeddingDto.getItemId()));
        metadata.put(RagConstant.EmbeddingMetaData.ITEM_ID, Util.null2String(embeddingDto.getItemId()));
        metadata.put(RagConstant.EmbeddingMetaData.COLLECTION_ID, Util.null2String(embeddingDto.getCollectionId()));
        metadata.put(RagConstant.EmbeddingMetaData.BASE_ID, Util.null2String(embeddingDto.getBaseId()));

        // TODO 库的id 表的id 支持自定义字段
        Document document = new Document(embeddingDto.getContent(), metadata);

        // 2. 文本切片 后续需要抽离出去
        // DocumentSplitter documentSplitter = DocumentSplitters.recursive(1000, 0, new
        // OpenAiTokenizer(GPT_3_5_TURBO));
        // DocumentSplitter documentSplitter = DocumentSplitters.recursive(384, 50);
        DocumentSplitter documentSplitter = DocumentSplitterUtil.getDocumentSplitter(
                embeddingDto.getMaxParagraphSize(),
                embeddingDto.getMaxOverlapSize(),
                embeddingDto.getSplitStrategy());

        // 3. 获取向量模型
        EmbeddingModel embeddingModel = vectorModelFactory.getVectorModel(embeddingDto.getVectorModel());

        // 4. 获取向量存储对象
        MilvusConfig config = milvusConfigService.getConfigByMark(RagConstant.MILVUS_MARK);
        EmbeddingStore<TextSegment> embeddingStore = langchainStoreBuilder.getOrCreateMilvusEmbeddingStore(config,
                embeddingDto.getCollectionIdentifier(), embeddingDto.getDimension());

        // 5. 向量存储 - 添加重试机制
        EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(documentSplitter)
                .textSegmentTransformer(embeddingDto.getTextSegmentTransformer() == null ? simpleTextSegmentTransformer
                        : embeddingDto.getTextSegmentTransformer())
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        
        // 执行向量存储，带重试机制
        executeWithRetry(() -> {
            embeddingStoreIngestor.ingest(document);
            return null;
        }, "向量存储");
        
        log.info("向量存储成功");
    }

    /**
     * 带重试机制的执行方法
     * @param operation 要执行的操作
     * @param operationName 操作名称（用于日志）
     * @return 操作结果
     * @throws Exception 异常
     */
    private <T> T executeWithRetry(RetryableOperation<T> operation, String operationName) throws Exception {
        Exception lastException = null;
        
        for (int attempt = 1; attempt <= milvusEmbeddingConfig.getMaxRetryAttempts(); attempt++) {
            try {
                return operation.execute();
            } catch (Exception e) {
                lastException = e;
                
                // 检查是否是速率限制错误
                if (isRateLimitError(e) && attempt < milvusEmbeddingConfig.getMaxRetryAttempts()) {
                    long delay = milvusEmbeddingConfig.getRetryDelayMs() * attempt; // 递增延迟
                    log.warn("{} 遇到速率限制，第 {} 次重试，延迟 {} ms", operationName, attempt, delay);
                    Thread.sleep(delay);
                    continue;
                }
                
                // 如果不是速率限制错误或已达到最大重试次数，直接抛出异常
                throw e;
            }
        }
        
        throw lastException;
    }
    
    /**
     * 判断是否为速率限制错误
     * @param e 异常
     * @return 是否为速率限制错误
     */
    private boolean isRateLimitError(Exception e) {
        String message = e.getMessage();
        return message != null && (
            message.contains("rate limit exceeded") ||
            message.contains("RateLimiter middleware") ||
            message.contains("error code: 8")
        );
    }
    
    /**
     * 可重试操作的函数式接口
     */
    @FunctionalInterface
    private interface RetryableOperation<T> {
        T execute() throws Exception;
    }

    /**
     * 删除知识条目的向量数据
     * 
     * @param EmbeddingDto
     * @return 是否删除成功
     */
    @Override
    public boolean deleteEmbedding(EmbeddingDto embeddingDto) {
        try {
            String itemId = Util.null2String(embeddingDto.getItemId());
            Filter filter = new MetadataFilterBuilder(RagConstant.EmbeddingMetaData.ITEM_ID).isEqualTo(itemId);

            MilvusConfig config = milvusConfigService.getConfigByMark(RagConstant.MILVUS_MARK);
            EmbeddingStore<TextSegment> embeddingStore = langchainStoreBuilder.getOrCreateMilvusEmbeddingStore(config,
                    embeddingDto.getCollectionIdentifier(), embeddingDto.getDimension());
            
            // 删除操作也添加重试机制
            executeWithRetry(() -> {
                embeddingStore.removeAll(filter);
                return null;
            }, "删除向量数据");
            
            return true;
        } catch (Exception e) {
            log.error("Failed to delete vector data for knowledge item: {}", embeddingDto.getItemId(), e);
            throw new RuntimeException("删除向量数据失败: " + e.getMessage(), e);
        }
    }
}
